#include "StdAfx.hpp"
#include "Video.hpp"
#include <Vfw.h>


const tchar * AVIFileError::GetMessage(HRESULT code)
{
	switch (code)
	{
	case AVIERR_OK:             return _T("OK");
	case AVIERR_UNSUPPORTED:    return _T("Unsupported.");
	case AVIERR_BADFORMAT:      return _T("Bad format");
	case AVIERR_MEMORY:         return _T("Memory");
	case AVIERR_INTERNAL:       return _T("Internal");
	case AVIERR_BADFLAGS:       return _T("Bad flags");
	case AVIERR_BADPARAM:       return _T("Bad param");
	case AVIERR_BADSIZE:        return _T("Bad size");
	case AVIERR_BADHANDLE:      return _T("Bad handle");
	case AVIERR_FILEREAD:       return _T("File read");
	case AVIERR_FILEWRITE:      return _T("File write");
	case AVIERR_FILEOPEN:       return _T("File open");
	case AVIERR_COMPRESSOR:     return _T("Compressor");
	case AVIERR_NOCOMPRESSOR:   return _T("No compressor");
	case AVIERR_READONLY:       return _T("Read only");
	case AVIERR_NODATA:         return _T("No data");
	case AVIERR_BUFFERTOOSMALL: return _T("Buffer too small");
	case AVIERR_CANTCOMPRESS:   return _T("Can't compress");
	case AVIERR_USERABORT:      return _T("User abort");
	case AVIERR_ERROR:          return _T("Error");
	default: return NULL;
	}
}

AVIFileError::AVIFileError(HRESULT hr, const tstring &Msg, const tstring &File, int Line)
{
	const tchar *codeMsg = GetMessage(hr);
	if (codeMsg == NULL)
		codeMsg = _T("Unknown");

	Push( _T("(AVIFileError,") + IntToStrR(hr) + _T(") ") + codeMsg);
	Push(Msg, File, Line);
}


const tchar * VCMError::GetMessage(HRESULT code)
{
	switch (code)
	{
	case ICERR_OK:           return _T("OK");
	case ICERR_DONTDRAW:     return _T("Don't draw");
	case ICERR_NEWPALETTE:   return _T("New palette");
	case ICERR_GOTOKEYFRAME: return _T("Got keyframe");
	case ICERR_STOPDRAWING:  return _T("Stop drawing");
	case ICERR_UNSUPPORTED:  return _T("Unsupported");
	case ICERR_BADFORMAT:    return _T("Bad format");
	case ICERR_MEMORY:       return _T("Memory");
	case ICERR_INTERNAL:     return _T("Internal");
	case ICERR_BADFLAGS:     return _T("Bad flags");
	case ICERR_BADPARAM:     return _T("Bad param");
	case ICERR_BADSIZE:      return _T("Bad size");
	case ICERR_BADHANDLE:    return _T("Bad handle");
	case ICERR_CANTUPDATE:   return _T("Can't update");
	case ICERR_ABORT:        return _T("Abort");
	case ICERR_ERROR:        return _T("Error");
	case ICERR_BADBITDEPTH:  return _T("Bad bit depth");
	case ICERR_BADIMAGESIZE: return _T("Bad image size");
	default: return NULL;
	}
}

VCMError::VCMError(HRESULT hr, const tstring &Msg, const tstring &File, int Line)
{
	const tchar *codeMsg = GetMessage(hr);
	if (codeMsg == NULL)
		codeMsg = _T("Unknown");

	Push( _T("(Video Compression Manager,") + IntToStrR(hr) + _T(") ") + codeMsg);
	Push(Msg, File, Line);
}


class VideoCompressorPimpl
{
public:
	COMPVARS CompVars;
	bool Compression;
	IAVIFile *AviFile;
	bool SeqCompressFrameStarted;
	HGLOBAL DstFormatHandle;
	IAVIStream *Stream;
	unsigned FrameIndex;
	size_t UncompressedFrameSize;
};


VideoCompressor::VideoCompressor()
: pimpl(new VideoCompressorPimpl)
{
	pimpl->Compression = false;
	pimpl->AviFile = NULL;
	pimpl->DstFormatHandle = NULL;
	pimpl->SeqCompressFrameStarted = false;
	pimpl->Stream = NULL;
	pimpl->FrameIndex = NULL;
	pimpl->UncompressedFrameSize = 0;

	AVIFileInit();
	ZeroMemory(&pimpl->CompVars, sizeof(pimpl->CompVars));
	pimpl->CompVars.cbSize = sizeof(pimpl->CompVars);
}

VideoCompressor::~VideoCompressor()
{
	EndFile();
	ICCompressorFree(&pimpl->CompVars);
	AVIFileExit();
}

bool VideoCompressor::ShowDialog(HWND parent, const char *title)
{
	BOOL ok = ICCompressorChoose(parent, 0, NULL, NULL, &pimpl->CompVars, (LPSTR)title);
	if (!ok)
		return false;
	pimpl->Compression = pimpl->CompVars.hic != NULL;
	return true;
}

void VideoCompressor::BeginFile(const char *filePath, const BITMAPINFOHEADER &srcFormat, unsigned fps, unsigned frameCount)
{
	assert(pimpl->AviFile == NULL && pimpl->DstFormatHandle == NULL && pimpl->SeqCompressFrameStarted == false && pimpl->Stream == NULL);
	pimpl->FrameIndex = 0;
	HRESULT hr;

	DWORD dstFormatSize;
	void *dstFormat;

	if (pimpl->Compression)
	{
		dstFormatSize = ICCompressGetFormatSize(pimpl->CompVars.hic, &srcFormat);
		if (dstFormatSize == 0 || dstFormatSize > 1024*1024)
			throw Error("Codec returned invalid format size.", __FILE__, __LINE__);
		pimpl->DstFormatHandle = GlobalAlloc(GHND, dstFormatSize);
		dstFormat = GlobalLock(pimpl->DstFormatHandle);
		hr = ICCompressGetFormat(pimpl->CompVars.hic, &srcFormat, dstFormat);
		if (FAILED(hr))
			throw VCMError(hr, "Cannot get compressed format.", __FILE__, __LINE__);

		hr = ICCompressQuery(pimpl->CompVars.hic, &srcFormat, dstFormat);
		if (FAILED(hr))
			throw VCMError(hr, "ICCompressQuery failed", __FILE__, __LINE__);
	}
	else
	{
		dstFormatSize = sizeof(BITMAPINFOHEADER);
		dstFormat = (LPVOID)&srcFormat;
		pimpl->UncompressedFrameSize = srcFormat.biWidth * srcFormat.biHeight * (srcFormat.biBitCount/8);
	}

	DeleteFile(filePath);

	hr = AVIFileOpen(&pimpl->AviFile, filePath, OF_CREATE | OF_WRITE, NULL);
	if (FAILED(hr))
		throw AVIFileError(hr, "Cannot open \"" + string(filePath) + "\"", __FILE__, __LINE__);

	unsigned rateScale = 1000;
	const char *streamName = "A";

	AVISTREAMINFO streamInfo;
	ZeroMemory(&streamInfo, sizeof(streamInfo));
	streamInfo.fccType = streamtypeVIDEO;
	streamInfo.fccHandler = pimpl->CompVars.fccHandler;
	streamInfo.dwScale = rateScale;
	streamInfo.dwRate = fps*rateScale;
	streamInfo.dwLength = frameCount;
	streamInfo.rcFrame.right  = srcFormat.biWidth;
	streamInfo.rcFrame.bottom = srcFormat.biHeight;
	streamInfo.dwSuggestedBufferSize = 0;
	strcpy_s(streamInfo.szName, _countof(streamInfo.szName), streamName);
	hr = AVIFileCreateStream(pimpl->AviFile, &pimpl->Stream, &streamInfo);
	if (FAILED(hr))
		throw AVIFileError(hr, "AVIFileCreateStream failed", __FILE__, __LINE__);

	hr = AVIStreamSetFormat(pimpl->Stream, 0, dstFormat, dstFormatSize);
	if (FAILED(hr))
		throw AVIFileError(hr, "AVIStreamSetFormat failed", __FILE__, __LINE__);

	if (pimpl->Compression)
	{
		BOOL b = ICSeqCompressFrameStart(&pimpl->CompVars, (BITMAPINFO*)&srcFormat);
		if (!b)
			throw Error("ICSeqCompressFrameStart failed");
		pimpl->SeqCompressFrameStarted = true;
	}
}

void VideoCompressor::EndFile()
{
	if (pimpl->SeqCompressFrameStarted)
	{
		ICSeqCompressFrameEnd(&pimpl->CompVars);
		pimpl->SeqCompressFrameStarted = false;
	}

	if (pimpl->DstFormatHandle)
	{
		GlobalUnlock(pimpl->DstFormatHandle);
		GlobalFree(pimpl->DstFormatHandle);
		pimpl->DstFormatHandle = NULL;
	}

	if (pimpl->Stream)
	{
		AVIStreamRelease(pimpl->Stream);
		pimpl->Stream = NULL;
	}

	if (pimpl->AviFile)
	{
		AVIFileRelease(pimpl->AviFile);
		pimpl->AviFile = NULL;
	}
}

void VideoCompressor::PostFrame(const void *data)
{
	assert(pimpl->Stream);

	if (pimpl->Compression)
	{
		BOOL wasKeyframe = FALSE;
		LONG thisFrameSize = 0x7fffffff;
		void *compressedData = ICSeqCompressFrame(&pimpl->CompVars, 0, (LPVOID)data, &wasKeyframe, &thisFrameSize);
		if (compressedData == NULL)
			throw Error("ICSeqCompressFrame returned NULL");

		HRESULT hr = AVIStreamWrite(pimpl->Stream, (LONG)pimpl->FrameIndex, 1, compressedData, thisFrameSize, wasKeyframe ? AVIIF_KEYFRAME : 0, NULL, NULL);
		if (FAILED(hr))
			throw AVIFileError(hr, "AVIStreamWrite failed", __FILE__, __LINE__);
	}
	else
	{
		HRESULT hr = AVIStreamWrite(pimpl->Stream, (LONG)pimpl->FrameIndex, 1, (LPVOID)data, pimpl->UncompressedFrameSize, AVIIF_KEYFRAME, NULL, NULL);
		if (FAILED(hr))
			throw AVIFileError(hr, "AVIStreamWrite failed", __FILE__, __LINE__);
	}

	pimpl->FrameIndex++;
}

void VideoCompressor::ResetCompression()
{
	pimpl->Compression = false;
}
